package distributed_rate_limiter

import (
	"Limiter/connect/rds"
	"context"
	"fmt"
	"github.com/redis/go-redis/v9"
	xrate "golang.org/x/time/rate"
	"log/slog"
	"strconv"
	"sync"
	"sync/atomic"
	"time"
)

type TokenBucketLimiter struct {
	// 每秒生产速率
	rate string
	// 桶容量
	cap string
	// 存储容器
	store *redis.Conn
	//脚本缓存标识
	sha string
	// 令牌桶 key
	tokenKey string
	// 缓存时间戳key
	timestampKey string
	// lock
	rescueLock sync.Mutex
	// redis健康标识
	redisAlive uint32
	// redis故障时采用进程内 令牌桶限流器
	rescueLimiter *xrate.Limiter
	// redis监控探测任务标识
	monitorStarted bool
}

var pingInterval = time.Second

var tokenBucketScript = `local rate=tonumber(ARGV[1]) local capacity=tonumber(ARGV[2]) local now = tonumber(ARGV[3]) local need = tonumber(ARGV[4]) local fill_time=capacity/rate local expire=math.floor(fill_time*2) local last_tokens=tonumber(redis.call("get",KEYS[1])) if last_tokens==nil then last_tokens=capacity end local last=tonumber(redis.call("get",KEYS[2])) if last==nil then last=0 end local delta=math.max(0,now-last) local now_tokens = math.min(capacity, last_tokens+(delta*rate)) local allowed =now_tokens>=need if allowed then now_tokens=now_tokens-need end redis.call("setex",KEYS[1],expire,now_tokens) redis.call("setex",KEYS[2],expire,now) return allowed`

func loadScript(ctx context.Context, client *redis.Client, script string) (sha string, err error) {
	s := redis.NewScript(script)
	sha, err = s.Load(ctx, client).Result()
	return
}

func NewTokenLimiter(ctx context.Context, rate int, cap int, client *redis.Client, key string) (tl *TokenBucketLimiter, err error) {
	var sha string
	tokenKey := fmt.Sprintf("tokenKey_%s", key)
	timestampKey := fmt.Sprintf("timestampKey_%s", tokenKey)

	sha, err = loadScript(ctx, client, tokenBucketScript)
	if err != nil {
		return
	}

	tl = &TokenBucketLimiter{
		rate:          strconv.Itoa(rate),
		cap:           strconv.Itoa(cap),
		store:         client.Conn(),
		tokenKey:      tokenKey,
		sha:           sha,
		timestampKey:  timestampKey,
		rescueLimiter: xrate.NewLimiter(xrate.Every(time.Second/time.Duration(rate)), cap),
		redisAlive:    1,
	}
	return
}

func (limiter *TokenBucketLimiter) Take(ctx context.Context, now time.Time, n int) bool {
	var err error
	var resp interface{}
	if atomic.LoadUint32(&limiter.redisAlive) == 0 {
		//开启兜底策略
		return limiter.rescueLimiter.AllowN(now, n)
	}
	resp, err = limiter.store.EvalSha(ctx, limiter.sha, []string{limiter.tokenKey, limiter.timestampKey}, limiter.rate, limiter.cap, strconv.FormatInt(now.Unix(), 10), n).Result()
	if err == redis.Nil {
		return false
	} else if err != nil {
		slog.Error(fmt.Sprintf("fail to use rate limiter: %s, use in-process limiter for rescue", err.Error()))
		limiter.startMonitor()
		return limiter.rescueLimiter.AllowN(now, n)
	}
	code, ok := resp.(int64)
	if !ok {
		slog.Error(fmt.Sprintf("fail to eval redis token_bucket_script: %v, use in-process limiter for rescue", resp))
		limiter.startMonitor()
		return limiter.rescueLimiter.AllowN(now, n)
	}
	return code == 1
}

// redis健康探测定时任务
func (limiter *TokenBucketLimiter) waitForRedis() {
	//定时器
	ticker := time.NewTicker(pingInterval)
	// 健康探测成功时回调此函数
	defer func() {
		ticker.Stop()
		limiter.rescueLock.Lock()
		limiter.monitorStarted = false
		limiter.rescueLock.Unlock()
	}()
	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(pingInterval))
	defer cancel()

	for range ticker.C {
		hashes, client, err := rds.ReLinkAndLoadScript(ctx, []string{tokenBucketScript})
		if err != nil {
			slog.Error(fmt.Sprintf("Error reloading scripts: %v", err))
			if client != nil {
				// 确保在发生错误时关闭连接
				client.Close()
			}
			continue
		}

		if len(hashes) < 1 {
			slog.Error("No hashes returned from Redis")
			if client != nil {
				// 确保在没有有效hashes时关闭连接
				client.Close()
			}
			continue
		}

		limiter.sha = hashes[0]
		limiter.store = client.Conn()
		atomic.StoreUint32(&limiter.redisAlive, 1)
		slog.Info(fmt.Sprintf("Redis returns to normal"))
		cancel()
		return
	}
}

// 开启redis健康探测
func (limiter *TokenBucketLimiter) startMonitor() {
	limiter.rescueLock.Lock()
	defer limiter.rescueLock.Unlock()
	// 防止重复开启
	if limiter.monitorStarted {
		return
	}
	// 设置任务和健康标识
	limiter.monitorStarted = true
	atomic.StoreUint32(&limiter.redisAlive, 0)
	// 健康探测
	go limiter.waitForRedis()
}
